package com.hys.ai.controller;

import lombok.SneakyThrows;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.InputStreamResource;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;

@RequestMapping("document")
@RestController
public class DocumentController {
    @Autowired
    private VectorStore vectorStore;

    /**
     * 嵌入文件
     *
     * @return 是否成功
     */
    @SneakyThrows
    @PostMapping("embedding")
    public Boolean embedding(List<MultipartFile> file) {
        for (MultipartFile multipartFile : file) {
            // 从IO流中读取文件
            TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(new InputStreamResource(multipartFile.getInputStream()));
            // 将文本内容划分成更小的块
            List<Document> splitDocuments = new TokenTextSplitter()
                    .apply(tikaDocumentReader.read());
            // 存入向量数据库，这个过程会自动调用embeddingModel,将文本变成向量再存入。
            vectorStore.add(splitDocuments);
        }
        return true;
    }

}
